package com.leetcode.design;

import java.util.HashSet;
import java.util.Set;

class Trie {

    private final TrieNode root;
    private final Set<String> strs;
    /**
     * Initialize your data structure here.
     */
    public Trie() {
        root = new TrieNode();
        strs = new HashSet<>();
    }

    /**
     * Inserts a word into the trie.
     */
    public void insert(String word) {
        strs.add(word);
        char[] array = word.toCharArray();
        TrieNode temp = root;
        for (char c : array) {
            if (!temp.containsKey(c)) {
                temp.put(c, new TrieNode());
            }
            temp = temp.get(c);
        }
        temp.setEnd();
    }

    /**
     * Returns if the word is in the trie.
     */
    public boolean search(String word) {
        return strs.contains(word);
    }

    /**
     * Returns if there is any word in the trie that starts with the given prefix.
     */
    public boolean startsWith(String prefix) {
        char[] array = prefix.toCharArray();
        TrieNode temp = root;
        for (char c : array) {
            if (temp.containsKey(c)) temp = temp.get(c);
            else return false;
        }
        return true;
    }

    static class TrieNode {

        // R links to node children
        private final TrieNode[] links;

        private final int R = 26;

        private boolean isEnd;

        public TrieNode() {
            links = new TrieNode[R];
        }

        public boolean containsKey(char ch) {
            return links[ch - 'a'] != null;
        }

        public TrieNode get(char ch) {
            return links[ch - 'a'];
        }

        public void put(char ch, TrieNode node) {
            links[ch - 'a'] = node;
        }

        public void setEnd() {
            isEnd = true;
        }

        public boolean isEnd() {
            return isEnd;
        }
    }

}
